package uestc.lj.basic.graph;

import java.util.Comparator;
import java.util.HashSet;
import java.util.PriorityQueue;
import java.util.Set;

/**
 * 最小生成树
 *
 * @Author:Crazlee
 * @Date:2021/11/24
 */
public class Code05_Prime {
	public static Set<Edge> primMST(Graph graph) {
		//解锁的边进入小根堆
		PriorityQueue<Edge> priorityQueue = new PriorityQueue<>(new Comparator<Edge>() {
			@Override
			public int compare(Edge o1, Edge o2) {
				return o1.weight - o2.weight;
			}
		});
		Set<Edge> result = new HashSet<>();

		//标记点是否被访问过
		HashSet<Node> set = new HashSet<>();

		for (Node node : graph.nodes.values()) {

			if (!set.contains(node)) {
				set.add(node);
				priorityQueue.addAll(node.edges);
				while (!priorityQueue.isEmpty()) {
					Edge edge = priorityQueue.poll();
					Node toNode = edge.to;
					if (!set.contains(toNode)) {
						set.add(toNode);
						result.add(edge);
						priorityQueue.addAll(toNode.edges);
					}
				}
			}

		}
		return result;
	}
}
